package leetcode.templates.unions;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;



public class Kruskal {


	class Edge {
		int x, y;
		int w;
		int i;

		public Edge(int x, int y, int w, int i) {
			this.x = x;
			this.y = y;
			this.w = w;
			this.i = i;
		}
	}


	public List<List<Integer>> findCriticalAndPseudoCriticalEdges(int n, int[][] edges) {
		int len = edges.length;
		// k算法要排序，为了保存index信息
		Edge[] myEdges = new Edge[len];
		for (int i = 0; i < len; i++) {
			int[] cur = edges[i];
			myEdges[i] = new Edge(cur[0], cur[1], cur[2], i);
		}
		Arrays.sort(myEdges, (a, b) -> a.w - b.w);
		int minWay = 0;// 最小生成树权值
		UnionFind union = new UnionFind(n);
		for (int i = 0; i < len; i++) {
			Edge edge = myEdges[i];
			if (!union.isSameSet(edge.x, edge.y)) {
				minWay += edge.w;
				union.union(edge.x, edge.y);
			}
		}
		// 关键边
		List<Integer> vips1 = new ArrayList<>();
		// 伪关键边
		List<Integer> vips2 = new ArrayList<>();
		for (int i = 0; i < len; i++) {
			Edge edge = myEdges[i];
			// false myEdge中第i条边一定不要
			int cur = this.weights(len, n, myEdges, i, false);
			if (cur > minWay) {// 说明这个边不能去掉，去掉权值上去了，是关建边。
				vips1.add(edge.i);
			} else if (cur == minWay) {
				// true myEdge中第i条边一定要
				cur = this.weights(len, n, myEdges, i, true);
				if (cur == minWay) {// 仍能完成构建最小生成树，说明这个边是可以去掉，也可以要，是个伪关键边
					vips2.add(edge.i);
				}
			}
		}
		List<List<Integer>> ans = new ArrayList<>();
		ans.add(vips1);
		ans.add(vips2);
		return ans;
	}

	class UnionFind {
		int[] roots;
		int[] sizes;
		int[] stack;

		public UnionFind(int n) {
			roots = new int[n];
			sizes = new int[n];
			for (int i = 0; i < n; i++) {
				roots[i] = i;
				sizes[i] = 1;
			}
			stack = new int[n];
		}

		public int findRoot(int id) {
			int root = 0;
			int p = 0;
			while ((root = roots[id]) != id) {
				stack[p++] = id;
				id = root;
			}
			while (p > 0) {
				roots[stack[--p]] = root;
			}
			return root;
		}

		public boolean isSameSet(int id1, int id2) {
			return findRoot(id1) == findRoot(id2);
		}

		public void union(int id1, int id2) {
			int p1 = this.findRoot(id1);
			int p2 = this.findRoot(id2);
			if (p1 != p2) {
				int size1 = this.sizes[p1];
				int size2 = this.sizes[p2];
				if (size1 < size2) {
					this.roots[p1] = p2;
					this.sizes[p2] = size1 + size2;
					this.sizes[p1] = 0;
				} else {
					this.roots[p2] = p1;
					this.sizes[p1] = size1 + size2;
					this.sizes[p2] = 0;
				}
			}
		}
	}

	private int weights(int len, int n, Edge[] myEdges, int idx, boolean add) {
		UnionFind union = new UnionFind(n);
		int minWay = 0;
		if (add) {
			Edge edge = myEdges[idx];
			union.union(edge.x, edge.y);
			minWay += edge.w;
			n--;
		}
		for (int i = 0; i < len; i++) {
			if (idx == i) {
				continue;
			}
			Edge edge = myEdges[i];
			if (!union.isSameSet(edge.x, edge.y)) {
				union.union(edge.x, edge.y);
				minWay += edge.w;
				n--;
			}
		}
		return n == 1 ? minWay : Integer.MAX_VALUE;
	}
}
